package ace.cmp.spring.graphql.core.execution.impl;

import ace.cmp.spring.graphql.core.execution.AceBatchLoaderRegistry;
import graphql.GraphQLContext;
import io.micrometer.context.ContextSnapshot;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.function.Supplier;
import org.dataloader.*;
import org.springframework.graphql.execution.DefaultBatchLoaderRegistry;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

/**
 * @author caspar
 * @date 2023/9/12 15:01 增加特性：
 * 1、支持断送DataLoader是否已经存在
 */
public class AceBatchLoaderRegistryImpl extends DefaultBatchLoaderRegistry
    implements AceBatchLoaderRegistry {
  private final List<ReactorBatchLoader<?, ?>> loaders = new ArrayList<>();

  private final List<ReactorMappedBatchLoader<?, ?>> mappedLoaders = new ArrayList<>();

  private final Supplier<DataLoaderOptions> defaultOptionsSupplier;

  /**
   * Default constructor
   */
  public AceBatchLoaderRegistryImpl() {
    this(DataLoaderOptions::newOptions);
  }

  /**
   * Constructor with a default {@link DataLoaderOptions} supplier to use as
   * a starting point for all registrations.
   *
   * @since 1.1.0
   */
  public AceBatchLoaderRegistryImpl(Supplier<DataLoaderOptions> defaultOptionsSupplier) {
    this.defaultOptionsSupplier = defaultOptionsSupplier;
  }

  @Override
  public <K, V> RegistrationSpec<K, V> forTypePair(Class<K> keyType, Class<V> valueType) {
    return new DefaultRegistrationSpec<>(valueType);
  }

  @Override
  public <K, V> RegistrationSpec<K, V> forName(String name) {
    return new DefaultRegistrationSpec<>(name);
  }

  @Override
  public void registerDataLoaders(DataLoaderRegistry registry, GraphQLContext context) {
    BatchLoaderContextProvider contextProvider = () -> context;
    DataLoaderOptions defaultOptions = this.defaultOptionsSupplier.get();
    for (ReactorBatchLoader<?, ?> loader : this.loaders) {
      DataLoaderOptions options = loader.getOptions();
      options =
          (options != null ? options : defaultOptions)
              .setBatchLoaderContextProvider(contextProvider);
      DataLoader<?, ?> dataLoader = DataLoaderFactory.newDataLoader(loader, options);
      registerDataLoader(loader.getName(), dataLoader, registry);
    }
    for (ReactorMappedBatchLoader<?, ?> loader : this.mappedLoaders) {
      DataLoaderOptions options = loader.getOptions();
      options =
          (options != null ? options : defaultOptions)
              .setBatchLoaderContextProvider(contextProvider);
      DataLoader<?, ?> dataLoader = DataLoaderFactory.newMappedDataLoader(loader, options);
      registerDataLoader(loader.getName(), dataLoader, registry);
    }
  }

  private void registerDataLoader(
      String name, DataLoader<?, ?> dataLoader, DataLoaderRegistry registry) {
    if (registry.getDataLoader(name) != null) {
      throw new IllegalStateException("More than one DataLoader named '" + name + "'");
    }
    registry.register(name, dataLoader);
  }

  private class DefaultRegistrationSpec<K, V> implements RegistrationSpec<K, V> {

    @Nullable private final Class<?> valueType;

    @Nullable private String name;

    @Nullable private DataLoaderOptions options;

    @Nullable private Consumer<DataLoaderOptions> optionsConsumer;

    public DefaultRegistrationSpec(Class<V> valueType) {
      this.valueType = valueType;
    }

    public DefaultRegistrationSpec(String name) {
      this.name = name;
      this.valueType = null;
    }

    @Override
    public RegistrationSpec<K, V> withName(String name) {
      this.name = name;
      return this;
    }

    @Override
    public RegistrationSpec<K, V> withOptions(Consumer<DataLoaderOptions> optionsConsumer) {
      this.optionsConsumer =
          (this.optionsConsumer != null
              ? this.optionsConsumer.andThen(optionsConsumer)
              : optionsConsumer);
      return this;
    }

    @Override
    public RegistrationSpec<K, V> withOptions(DataLoaderOptions options) {
      this.options = options;
      return this;
    }

    @Override
    public void registerBatchLoader(BiFunction<List<K>, BatchLoaderEnvironment, Flux<V>> loader) {
      AceBatchLoaderRegistryImpl.this.loaders.add(
          new ReactorBatchLoader<>(initName(), loader, initOptionsSupplier()));
    }

    @Override
    public void registerMappedBatchLoader(
        BiFunction<Set<K>, BatchLoaderEnvironment, Mono<Map<K, V>>> loader) {
      AceBatchLoaderRegistryImpl.this.mappedLoaders.add(
          new ReactorMappedBatchLoader<>(initName(), loader, initOptionsSupplier()));
    }

    @Nullable
    private Supplier<DataLoaderOptions> initOptionsSupplier() {
      if (this.options == null && this.optionsConsumer == null) {
        return null;
      }

      Supplier<DataLoaderOptions> optionsSupplier =
          (this.options != null ? () -> this.options : defaultOptionsSupplier);

      if (this.optionsConsumer == null) {
        return optionsSupplier;
      }

      return () -> {
        DataLoaderOptions options = optionsSupplier.get();
        this.optionsConsumer.accept(options);
        return options;
      };
    }

    private String initName() {
      if (StringUtils.hasText(this.name)) {
        return this.name;
      }
      Assert.notNull(
          this.valueType, "Value type not available to select a default DataLoader name.");
      return (StringUtils.hasText(this.name) ? this.name : this.valueType.getName());
    }
  }

  /**
   * {@link BatchLoaderWithContext} that delegates to a {@link Flux} batch
   * loading function and exposes Reactor context to it.
   */
  private static class ReactorBatchLoader<K, V> implements BatchLoaderWithContext<K, V> {

    private final String name;

    private final BiFunction<List<K>, BatchLoaderEnvironment, Flux<V>> loader;

    @Nullable private final Supplier<DataLoaderOptions> optionsSupplier;

    private ReactorBatchLoader(
        String name,
        BiFunction<List<K>, BatchLoaderEnvironment, Flux<V>> loader,
        @Nullable Supplier<DataLoaderOptions> optionsSupplier) {

      this.name = name;
      this.loader = loader;
      this.optionsSupplier = optionsSupplier;
    }

    public String getName() {
      return this.name;
    }

    @Nullable
    public DataLoaderOptions getOptions() {
      return (this.optionsSupplier != null ? this.optionsSupplier.get() : null);
    }

    @Override
    @SuppressWarnings("deprecation")
    public CompletionStage<List<V>> load(List<K> keys, BatchLoaderEnvironment environment) {
      GraphQLContext graphQLContext = environment.getContext();
      ContextSnapshot snapshot = ContextSnapshot.captureFrom(graphQLContext);
      try {
        return snapshot
            .wrap(
                () ->
                    this.loader
                        .apply(keys, environment)
                        .collectList()
                        .contextWrite(snapshot::updateContext)
                        .toFuture())
            .call();
      } catch (Exception ex) {
        return CompletableFuture.failedFuture(ex);
      }
    }
  }

  /**
   * {@link MappedBatchLoaderWithContext} that delegates to a {@link Mono}
   * batch loading function and exposes Reactor context to it.
   */
  private static class ReactorMappedBatchLoader<K, V>
      implements MappedBatchLoaderWithContext<K, V> {

    private final String name;

    private final BiFunction<Set<K>, BatchLoaderEnvironment, Mono<Map<K, V>>> loader;

    @Nullable private final Supplier<DataLoaderOptions> optionsSupplier;

    private ReactorMappedBatchLoader(
        String name,
        BiFunction<Set<K>, BatchLoaderEnvironment, Mono<Map<K, V>>> loader,
        @Nullable Supplier<DataLoaderOptions> optionsSupplier) {

      this.name = name;
      this.loader = loader;
      this.optionsSupplier = optionsSupplier;
    }

    public String getName() {
      return this.name;
    }

    @Nullable
    public DataLoaderOptions getOptions() {
      return (this.optionsSupplier != null ? this.optionsSupplier.get() : null);
    }

    @Override
    @SuppressWarnings("deprecation")
    public CompletionStage<Map<K, V>> load(Set<K> keys, BatchLoaderEnvironment environment) {
      GraphQLContext graphQLContext = environment.getContext();
      ContextSnapshot snapshot = ContextSnapshot.captureFrom(graphQLContext);
      try {
        return snapshot
            .wrap(
                () ->
                    this.loader
                        .apply(keys, environment)
                        .contextWrite(snapshot::updateContext)
                        .toFuture())
            .call();
      } catch (Exception ex) {
        return CompletableFuture.failedFuture(ex);
      }
    }
  }

  @Override
  public Boolean hasMapperLoader(String name) {
    return this.mappedLoaders.stream().anyMatch(p -> p.getName().equals(name));
  }

  @Override
  public Boolean hasLoader(String name) {
    return this.loaders.stream().anyMatch(p -> p.getName().equals(name));
  }
}
